import numpy as np
from .base_optimizer import BaseOptimizer

class GA(BaseOptimizer):
    def __init__(self, dim, pop_size=50, max_iter=1000, crossover_rate=0.8, mutation_rate=0.1):
        super().__init__(dim, pop_size, max_iter)
        self.crossover_rate = crossover_rate
        self.mutation_rate = mutation_rate
        
    def select_parents(self, population, fitness):
        """轮盘赌选择"""
        fitness = 1 / (1 + fitness)  # 转换为最大化问题
        probs = fitness / fitness.sum()
        selected_idx = np.random.choice(len(population), size=2, p=probs)
        return population[selected_idx]
        
    def crossover(self, parent1, parent2):
        """模拟二进制交叉"""
        if np.random.random() < self.crossover_rate:
            mask = np.random.random(self.dim) < 0.5
            child1 = np.where(mask, parent1, parent2)
            child2 = np.where(mask, parent2, parent1)
            return child1, child2
        return parent1.copy(), parent2.copy()
        
    def mutate(self, individual, bounds):
        """高斯变异"""
        mask = np.random.random(self.dim) < self.mutation_rate
        mutation = np.random.normal(0, 0.1, self.dim)
        individual[mask] += mutation[mask]
        return self.clip_to_bounds(individual, bounds)
        
    def optimize(self, objective_func, bounds=(-5, 5)):
        # 初始化种群
        population = self.initialize_population(bounds)
        fitness = np.array([objective_func(ind) for ind in population])
        
        # 记录最优解
        best_idx = np.argmin(fitness)
        self.best_solution = population[best_idx].copy()
        self.best_fitness = fitness[best_idx]
        
        for gen in range(self.max_iter):
            new_population = []
            
            # 精英保留
            elite_idx = np.argmin(fitness)
            new_population.append(population[elite_idx].copy())
            
            # 生成新一代
            while len(new_population) < self.pop_size:
                # 选择父代
                parents = self.select_parents(population, fitness)
                
                # 交叉
                child1, child2 = self.crossover(parents[0], parents[1])
                
                # 变异
                child1 = self.mutate(child1, bounds)
                child2 = self.mutate(child2, bounds)
                
                new_population.extend([child1, child2])
            
            # 更新种群
            population = np.array(new_population[:self.pop_size])
            fitness = np.array([objective_func(ind) for ind in population])
            
            # 更新最优解
            best_idx = np.argmin(fitness)
            if fitness[best_idx] < self.best_fitness:
                self.best_solution = population[best_idx].copy()
                self.best_fitness = fitness[best_idx]
                
        return self.best_solution, self.best_fitness 